#pragma once
#include "httplib.h"
#include "logger.h"
#include <array>
#include <atomic>
#include <condition_variable>
#include <functional>
#include <future>
#include <memory>
#include <mutex>
#include <queue>
#include <string>
#include <thread>
#include <vector>

// 任务队列
enum TaskPriority
{
    HIGH,   // 音视频任务
    MEDIUM, // 消息相关任务
    LOW     // 好友请求和文件上传下载任务
};

// 任务类型
class Task
{
public:
    using TaskFunction = std::function<void()>;

    Task(TaskFunction func, TaskPriority priority = TaskPriority::MEDIUM,
         const std::string &description = "")
        : _function(std::move(func)), _priority(priority), _description(description) {}

    // 执行任务
    void execute()
    {
        if (_function)
        {
            _function();
        }
    }

    // 获取任务优先级
    TaskPriority getPriority() const
    {
        return _priority;
    }

    // 获取任务描述
    const std::string &getDescription() const
    {
        return _description;
    }

private:
    TaskFunction _function;   // 任务函数
    TaskPriority _priority;   // 任务优先级
    std::string _description; // 任务描述
};

// 自定义比较函数，用于优先级队列
// 注意：这里priority_queue默认是最大堆，我们希望高优先级在队列前面，所以需要反向比较
struct TaskPriorityCompare
{
    bool operator()(const std::shared_ptr<Task> &lhs, const std::shared_ptr<Task> &rhs) const
    {
        return lhs->getPriority() > rhs->getPriority(); // 数值小的优先级高
    }
};

// 优先级任务队列类
class TaskQueue
{
public:
    TaskQueue()
        : _stopping(false) {}
    ~TaskQueue()
    {
        stop();
    }

    // 停止队列
    void stop()
    {
        _stopping = true;
        _condition.notify_all();
    }

    // 添加任务到队列
    void pushTask(std::shared_ptr<Task> task)
    {
        std::unique_lock<std::mutex> lock(_mutex);

        // 根据任务优先级添加到相应队列
        switch (task->getPriority())
        {
        case TaskPriority::HIGH:
            _highPriorityQueue.push(task);
            break;
        case TaskPriority::MEDIUM:
            _mediumPriorityQueue.push(task);
            break;
        case TaskPriority::LOW:
            _lowPriorityQueue.push(task);
            break;
        }

        // 通知一个等待的线程有新任务
        _condition.notify_one();
    }

    // 获取任务，优先返回高优先级任务
    std::shared_ptr<Task> popTask()
    {
        std::unique_lock<std::mutex> lock(_mutex);

        // 等待直到有任务或队列停止
        _condition.wait(lock, [this]
                        { return _stopping ||
                                 !_highPriorityQueue.empty() ||
                                 !_mediumPriorityQueue.empty() ||
                                 !_lowPriorityQueue.empty(); });

        // 如果队列停止且没有任务，返回nullptr
        if (_stopping &&
            _highPriorityQueue.empty() &&
            _mediumPriorityQueue.empty() &&
            _lowPriorityQueue.empty())
        {
            return nullptr;
        }

        // 按优先级返回任务
        std::shared_ptr<Task> task;
        if (!_highPriorityQueue.empty())
        {
            task = _highPriorityQueue.front();
            _highPriorityQueue.pop();
        }
        else if (!_mediumPriorityQueue.empty())
        {
            task = _mediumPriorityQueue.front();
            _mediumPriorityQueue.pop();
        }
        else if (!_lowPriorityQueue.empty())
        {
            task = _lowPriorityQueue.front();
            _lowPriorityQueue.pop();
        }

        return task;
    }

    // 获取任务，不阻塞，如果没有任务则返回nullptr
    std::shared_ptr<Task> tryPopTask()
    {
        std::unique_lock<std::mutex> lock(_mutex, std::try_to_lock);
        if (!lock.owns_lock())
        {
            return nullptr;
        }

        if (_highPriorityQueue.empty() &&
            _mediumPriorityQueue.empty() &&
            _lowPriorityQueue.empty())
        {
            return nullptr;
        }

        // 按优先级返回任务
        std::shared_ptr<Task> task;
        if (!_highPriorityQueue.empty())
        {
            task = _highPriorityQueue.front();
            _highPriorityQueue.pop();
        }
        else if (!_mediumPriorityQueue.empty())
        {
            task = _mediumPriorityQueue.front();
            _mediumPriorityQueue.pop();
        }
        else if (!_lowPriorityQueue.empty())
        {
            task = _lowPriorityQueue.front();
            _lowPriorityQueue.pop();
        }

        return task;
    }

    // 获取特定优先级的任务数量
    size_t getTaskCount(TaskPriority priority) const
    {
        std::unique_lock<std::mutex> lock(_mutex);
        switch (priority)
        {
        case TaskPriority::HIGH:
            return _highPriorityQueue.size();
        case TaskPriority::MEDIUM:
            return _mediumPriorityQueue.size();
        case TaskPriority::LOW:
            return _lowPriorityQueue.size();
        default:
            return 0;
        }
    }

    // 获取总任务数量
    size_t getTotalTaskCount() const
    {
        std::unique_lock<std::mutex> lock(_mutex);
        return _highPriorityQueue.size() +
               _mediumPriorityQueue.size() +
               _lowPriorityQueue.size();
    }

    // 检查队列是否为空
    bool isEmpty() const
    {
        std::unique_lock<std::mutex> lock(_mutex);
        return _highPriorityQueue.empty() &&
               _mediumPriorityQueue.empty() &&
               _lowPriorityQueue.empty();
    }

private:
    std::queue<std::shared_ptr<Task>> _highPriorityQueue;   // 高优先级队列
    std::queue<std::shared_ptr<Task>> _mediumPriorityQueue; // 中优先级队列
    std::queue<std::shared_ptr<Task>> _lowPriorityQueue;    // 低优先级队列

    mutable std::mutex _mutex;          // 互斥锁
    std::condition_variable _condition; // 条件变量
    std::atomic<bool> _stopping;        // 停止标志
};

// 基本线程池类
class ThreadPool
{
public:
    // 构造函数
    ThreadPool(size_t numThreads, const std::string &poolName = "")
        : _poolName(poolName), _running(true), _numThreads(numThreads)
    {
        startThreads(numThreads);
    }

    // 析构函数
    virtual ~ThreadPool()
    {
        shutdown();
    }

    // 提交任务到线程池并获取future
    template <typename F, typename... Args>
    auto submitTask(TaskPriority priority, F &&f, Args &&...args)
    {
        using ReturnType = std::invoke_result_t<F, Args...>;
        auto taskPackage = std::make_shared<std::packaged_task<ReturnType()>>(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...));

        auto task = std::make_shared<Task>(
            [taskPackage]()
            { (*taskPackage)(); },
            priority,
            _poolName + " Task");

        _taskQueue.pushTask(task);
        return taskPackage->get_future();
    }

    // 向线程池提交任务，无返回值
    void enqueueTask(std::shared_ptr<Task> task)
    {
        _taskQueue.pushTask(task);
    }

    // 调整线程池大小
    void resize(size_t numThreads)
    {
        if (numThreads == _threads.size())
        {
            return;
        }

        if (numThreads > _threads.size())
        {
            // 增加线程
            size_t threadsToAdd = numThreads - _threads.size();
            startThreads(threadsToAdd);
        }
        else
        {
            // 减少线程 (需要实现减少线程的逻辑)
            // 这里暂不实现
        }
    }

    // 获取线程池名称
    const std::string &getName() const
    {
        return _poolName;
    }

    // 获取线程池中线程数量
    size_t getThreadCount() const
    {
        return _threads.size();
    }

    // 获取任务队列中任务数量
    size_t getTaskCount() const
    {
        return _taskQueue.getTotalTaskCount();
    }

    // 检查线程池是否在运行
    bool isRunning() const
    {
        return _running;
    }

    // 检查线程池是否可以接收特定优先级的任务
    // 这里假设每个线程可以处理一个任务，所以如果空闲线程数量足够，就可以接收任务
    bool canAcceptTask(TaskPriority priority) const
    {
        // 简单实现，如果有空闲线程就可以接收任务
        size_t busyThreads = _busyThreads.load();
        return busyThreads < _threads.size();
    }

protected:
    // 启动线程
    void startThreads(size_t numThreads)
    {
        for (size_t i = 0; i < numThreads; i++)
        {
            _threads.emplace_back([this]
                                  { workerThread(); });
        }
    }

    // 工作线程函数
    void workerThread()
    {
        while (_running)
        {
            std::shared_ptr<Task> task = _taskQueue.popTask();

            if (task)
            {
                _busyThreads++;
                try
                {
                    task->execute();
                }
                catch (const std::exception &e)
                {
                    std::cerr << "Exception in thread pool task: " << e.what() << std::endl;
                }
                catch (...)
                {
                    std::cerr << "Unknown exception in thread pool task" << std::endl;
                }
                _busyThreads--;
            }
        }
    }

    // 关闭线程池
    void shutdown()
    {
        _running = false;
        _taskQueue.stop();

        // 等待所有线程结束
        for (auto &thread : _threads)
        {
            if (thread.joinable())
            {
                thread.join();
            }
        }

        _threads.clear();
    }

protected:
    std::string _poolName;               // 线程池名称
    std::atomic<bool> _running;          // 运行标志
    std::vector<std::thread> _threads;   // 线程容器
    std::atomic<size_t> _busyThreads{0}; // 忙碌线程数量
    TaskQueue _taskQueue;                // 任务队列
    size_t _numThreads;                  // 线程数量
};

// 多级线程池控制器
class ThreadPoolCtrl
{
public:
    // 单例模式获取实例
    static ThreadPoolCtrl &getInstance()
    {
        static ThreadPoolCtrl instance;
        return instance;
    }

    // 初始化线程池控制器
    void initialize(size_t highThreads = 4,
                    size_t mediumThreads = 6,
                    size_t lowThreads = 8)
    {
        std::lock_guard<std::mutex> lock(_mutex);

        // 创建三种不同线程池
        _pools[static_cast<size_t>(TaskPriority::HIGH)] =
            std::make_unique<ThreadPool>(highThreads, "High-Pool");

        _pools[static_cast<size_t>(TaskPriority::MEDIUM)] =
            std::make_unique<ThreadPool>(mediumThreads, "Medium-Pool");

        _pools[static_cast<size_t>(TaskPriority::LOW)] =
            std::make_unique<ThreadPool>(lowThreads, "Low-Pool");
    }

    // 提交任务到线程池控制器
    template <typename F, typename... Args>
    auto submitTask(TaskPriority priority, F &&f, Args &&...args)
    {
        // 根据任务优先级选择合适的线程池
        TaskPriority poolType = selectPoolForTask(priority);

        // 获取选择的线程池
        auto &pool = _pools[static_cast<size_t>(poolType)];

        // 提交任务到选择的线程池
        return pool->submitTask(priority, std::forward<F>(f), std::forward<Args>(args)...);
    }

    // 调整特定线程池的大小
    void resizePool(TaskPriority poolType, size_t numThreads)
    {
        std::lock_guard<std::mutex> lock(_mutex);

        if (_pools[static_cast<size_t>(poolType)])
        {
            _pools[static_cast<size_t>(poolType)]->resize(numThreads);
        }
    }

    // 获取特定线程池的线程数量
    size_t getPoolThreadCount(TaskPriority poolType) const
    {
        std::lock_guard<std::mutex> lock(_mutex);

        if (_pools[static_cast<size_t>(poolType)])
        {
            return _pools[static_cast<size_t>(poolType)]->getThreadCount();
        }

        return 0;
    }

    // 获取特定线程池的任务数量
    size_t getPoolTaskCount(TaskPriority poolType) const
    {
        std::lock_guard<std::mutex> lock(_mutex);

        if (_pools[static_cast<size_t>(poolType)])
        {
            return _pools[static_cast<size_t>(poolType)]->getTaskCount();
        }

        return 0;
    }

private:
    // 根据任务优先级和线程池负载选择合适的线程池
    TaskPriority selectPoolForTask(TaskPriority priority) const
    {
        std::lock_guard<std::mutex> lock(_mutex);

        // 默认映射关系
        TaskPriority defaultPool = priority;

        // 检查默认线程池是否可以接受该任务
        if (_pools[static_cast<size_t>(defaultPool)]->canAcceptTask(priority))
        {
            return defaultPool;
        }

        // 如果默认线程池负载高，尝试其它线程池
        if (priority == TaskPriority::HIGH)
        {
            // 音视频任务优先级最高，先尝试消息线程池，再尝试通用线程池
            if (_pools[static_cast<size_t>(TaskPriority::MEDIUM)]->canAcceptTask(priority))
            {
                return TaskPriority::MEDIUM;
            }

            // 如果消息线程池也负载高，则使用通用线程池
            // 如果通用线程池还不够，可以考虑扩容
            if (_pools[static_cast<size_t>(TaskPriority::LOW)]->canAcceptTask(priority))
            {
                return TaskPriority::LOW;
            }
            else
            {
                // 扩容通用线程池
                size_t currentThreads = _pools[static_cast<size_t>(TaskPriority::LOW)]->getThreadCount();
                _pools[static_cast<size_t>(TaskPriority::LOW)]->resize(currentThreads + 2);
                return TaskPriority::LOW;
            }
        }
        else if (priority == TaskPriority::MEDIUM)
        {
            // 消息任务优先尝试通用线程池
            if (_pools[static_cast<size_t>(TaskPriority::LOW)]->canAcceptTask(priority))
            {
                return TaskPriority::LOW;
            }
            else
            {
                // 扩容通用线程池
                size_t currentThreads = _pools[static_cast<size_t>(TaskPriority::LOW)]->getThreadCount();
                _pools[static_cast<size_t>(TaskPriority::LOW)]->resize(currentThreads + 2);
                return TaskPriority::LOW;
            }
        }

        // 对于LOW优先级任务，如果通用线程池负载高，则扩容
        size_t currentThreads = _pools[static_cast<size_t>(TaskPriority::LOW)]->getThreadCount();
        _pools[static_cast<size_t>(TaskPriority::LOW)]->resize(currentThreads + 2);
        return TaskPriority::LOW;
    }

    // 构造函数设为私有，确保单例模式
    ThreadPoolCtrl() {}

    // 禁止拷贝和赋值
    ThreadPoolCtrl(const ThreadPoolCtrl &) = delete;
    ThreadPoolCtrl &operator=(const ThreadPoolCtrl &) = delete;

private:
    // 三个不同类型的线程池
    std::array<std::unique_ptr<ThreadPool>, 3> _pools;

    // 互斥锁，用于保护对线程池的访问
    mutable std::mutex _mutex;
};